import os
import json
import shutil

# 定义类别和对应的group_id
category_to_group_id = {
    "box": 0,
    "arrows_right": 1,
    "arrows_left": 2
}

def merge_annotations(folder_path1, folder_path2, output_folder):
    common_files = set()

    for filename in os.listdir(folder_path1):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            continue

        annotation1_path = os.path.join(folder_path1, filename)
        annotation2_path = os.path.join(folder_path2, filename)

        if not os.path.exists(annotation2_path):
            with open(annotation1_path, 'r') as file1:
                print(annotation1_path)
                annotation1 = json.load(file1)

                wrong_name_flag = False
                if annotation1["shapes"][0]["label"] == "box0":
                    annotation1["shapes"][0]["label"] = "box"
                    wrong_name_flag = True

                # annotation1["shapes"][0]["group_id"] = category_to_group_id[annotation1["shapes"][0]["label"]]
                # for category in category_to_group_id:
                #     if annotation1["shapes"][0]["label"] == category:
                #         group_id = category_to_group_id[category]
                #         num_keypoints = len(annotation1["shapes"]) - 1
                #         for i in range(1, num_keypoints + 1):
                #             annotation1["shapes"][i]["group_id"] = group_id

                if wrong_name_flag:
                    annotation1["shapes"][1]["label"] = "box_1"
                    annotation1["shapes"][2]["label"] = "box_2"
                    annotation1["shapes"][3]["label"] = "box_3"
                    annotation1["shapes"][4]["label"] = "box_4"
                    annotation1["shapes"][5]["label"] = "box_5"
                    annotation1["shapes"][6]["label"] = "box_6"

                output_path = os.path.join(output_folder, filename)
                with open(output_path, 'w') as outfile:
                    json.dump(annotation1, outfile, indent=4)
        else:
            common_files.add(filename)

            with open(annotation1_path, 'r') as file1, open(annotation2_path, 'r') as file2:
                annotation1 = json.load(file1)
                annotation2 = json.load(file2)

                wrong_name_flag = 0
                if annotation1["shapes"][0]["label"] == "box0":
                    annotation1["shapes"][0]["label"] = "box"
                    wrong_name_flag = 1
                if annotation2["shapes"][0]["label"] == "box0":
                    annotation2["shapes"][0]["label"] = "box"
                    wrong_name_flag = 2

                # 填充group_id
                # annotation1["shapes"][0]["group_id"] = category_to_group_id[annotation1["shapes"][0]["label"]]
                # annotation2["shapes"][0]["group_id"] = category_to_group_id[annotation2["shapes"][0]["label"]]
                # for category in category_to_group_id:
                #     if annotation1["shapes"][0]["label"] == category:
                #         group_id = category_to_group_id[category]
                #         num_keypoints = len(annotation1["shapes"]) - 1
                #         for i in range(1, num_keypoints + 1):
                #             annotation1["shapes"][i]["group_id"] = group_id
                #             annotation2["shapes"][i]["group_id"] = group_id

                if wrong_name_flag == 1:
                    annotation1["shapes"][1]["label"] = "box_1"
                    annotation1["shapes"][2]["label"] = "box_2"
                    annotation1["shapes"][3]["label"] = "box_3"
                    annotation1["shapes"][4]["label"] = "box_4"
                    annotation1["shapes"][5]["label"] = "box_5"
                    annotation1["shapes"][6]["label"] = "box_6"
                if wrong_name_flag == 2:
                    annotation2["shapes"][1]["label"] = "box_1"
                    annotation2["shapes"][2]["label"] = "box_2"
                    annotation2["shapes"][3]["label"] = "box_3"
                    annotation2["shapes"][4]["label"] = "box_4"
                    annotation2["shapes"][5]["label"] = "box_5"
                    annotation2["shapes"][6]["label"] = "box_6"

                merged_shapes = annotation1["shapes"] + annotation2["shapes"]
                combined_annotation = {
                    "version": annotation1["version"],
                    "flags": annotation1["flags"],
                    "shapes": merged_shapes,
                    "imagePath": annotation1["imagePath"],
                    "imageData": annotation1["imageData"],
                    "imageHeight": annotation1["imageHeight"],
                    "imageWidth": annotation1["imageWidth"]
                }

                output_path = os.path.join(output_folder, filename)
                with open(output_path, 'w') as outfile:
                    json.dump(combined_annotation, outfile, indent=4)

    for filename in os.listdir(folder_path2):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            continue
        if filename not in common_files:
            with open(os.path.join(folder_path2, filename), 'r') as file2:
                annotation2 = json.load(file2)

                wrong_name_flag = False
                if annotation2["shapes"][0]["label"] == "box0":
                    annotation2["shapes"][0]["label"] = "box"
                    wrong_name_flag = True

                # annotation2["shapes"][0]["group_id"] = category_to_group_id[annotation2["shapes"][0]["label"]]
                # for category in category_to_group_id:
                #     if annotation2["shapes"][0]["label"] == category:
                #         group_id = category_to_group_id[category]
                #         num_keypoints = len(annotation2["shapes"]) - 1
                #         for i in range(1, num_keypoints + 1):
                #             annotation2["shapes"][i]["group_id"] = group_id

                if wrong_name_flag:
                    annotation2["shapes"][1]["label"] = "box_1"
                    annotation2["shapes"][2]["label"] = "box_2"
                    annotation2["shapes"][3]["label"] = "box_3"
                    annotation2["shapes"][4]["label"] = "box_4"
                    annotation2["shapes"][5]["label"] = "box_5"
                    annotation2["shapes"][6]["label"] = "box_6"

                output_path = os.path.join(output_folder, filename)
                with open(output_path, 'w') as outfile:
                    json.dump(annotation2, outfile, indent=4)

    print("Completed merging annotations.")




def copy_common_images(folder1, folder2, target_folder):
    # 获取文件夹1中的所有图片文件
    images_folder1 = [file for file in os.listdir(folder1) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]

    # 获取文件夹2中的所有图片文件
    images_folder2 = [file for file in os.listdir(folder2) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]

    # 找到两个文件夹中相同的图片文件
    common_images = set(images_folder1) & set(images_folder2)

    # 复制相同的图片文件到目标文件夹
    for image_file in common_images:
        source_path1 = os.path.join(folder1, image_file)
        source_path2 = os.path.join(folder2, image_file)
        target_path = os.path.join(target_folder, image_file)

        # 复制文件
        shutil.copyfile(source_path1, target_path)
        print(f"Copied {image_file} to {target_folder}")


# 指定两个存储标注文件的文件夹路径
folder_path1 = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/侧面正面完整/左/亮/红/arrows json_annotations"
folder_path2 = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/侧面正面完整/左/亮/红/box json_annotations"

output_folder = os.path.join(os.path.dirname(folder_path1), "merged_annotations")  # os.path.dirname得到上一级目录
os.makedirs(output_folder, exist_ok=True)

merge_annotations(folder_path1, folder_path2, output_folder)

copy_common_images(folder_path1, folder_path2, output_folder)